# -*- encoding: utf-8 -*-
import audioop
import os
import time
import websockets, ssl
import asyncio
# import threading
import argparse
import json
import traceback
from multiprocessing import Process
# from funasr.fileio.datadir_writer import DatadirWriter
import wave
import logging

logging.basicConfig(level=logging.ERROR)

parser = argparse.ArgumentParser()
parser.add_argument("--host",
                    type=str,
                    default="localhost",
                    required=False,
                    help="host ip, localhost, 0.0.0.0")
parser.add_argument("--port",
                    type=int,
                    default=10095,
                    required=False,
                    help="grpc server port")
parser.add_argument("--chunk_size",
                    type=str,
                    default="5, 10, 5",
                    help="chunk")
parser.add_argument("--encoder_chunk_look_back",
                    type=int,
                    default=4,
                    help="chunk")
parser.add_argument("--decoder_chunk_look_back",
                    type=int,
                    default=0,
                    help="chunk")
parser.add_argument("--chunk_interval",
                    type=int,
                    default=10,
                    help="chunk")
parser.add_argument("--hotword",
                    type=str,
                    default="",
                    help="hotword file path, one hotword perline (e.g.:阿里巴巴 20)")
parser.add_argument("--audio_in",
                    type=str,
                    default=None,
                    help="audio_in")
parser.add_argument("--audio_fs",
                    type=int,
                    default=16000,
                    help="audio_fs")
parser.add_argument("--send_without_sleep",
                    action="store_true",
                    default=True,
                    help="if audio_in is set, send_without_sleep")
parser.add_argument("--thread_num",
                    type=int,
                    default=1,
                    help="thread_num")
parser.add_argument("--words_max_print",
                    type=int,
                    default=10000,
                    help="chunk")
parser.add_argument("--output_dir",
                    type=str,
                    default=None,
                    help="output_dir")
parser.add_argument("--ssl",
                    type=int,
                    default=1,
                    help="1 for ssl connect, 0 for no ssl")
parser.add_argument("--use_itn",
                    type=int,
                    default=1,
                    help="1 for using itn, 0 for not itn")
parser.add_argument("--mode",
                    type=str,
                    default="2pass",
                    help="offline, online, 2pass")

args = parser.parse_args()
args.chunk_size = [int(x) for x in args.chunk_size.split(",")]
print(args)
# voices = asyncio.Queue()
from queue import Queue

voices = Queue()
offline_msg_done=False

if args.output_dir is not None:
    # if os.path.exists(args.output_dir):
    #     os.remove(args.output_dir)

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)


async def record_fs():
    global voices
    await websocket.send('{"command":"START","config":{"fs":8000,"format":"pcm","itn":false,"punc":false,"speaker":false,"vadTail":600,"maxSeconds":60,"save":true,"hotwordId":1}}')

    with open('/home/leo/Downloads/4357f7ee-a784-4684-a481-6d8bf1eda4bc.wav',
                   'rb') as file:
        speech = file.read()
        file.close()
        speech = audioop.alaw2lin(speech,2)

        speech_length = len(speech)
        sample_offset = 0
        stride_size = 320 * 10

        is_final = False
        for sample_offset in range(0, speech_length, min(stride_size, speech_length - sample_offset)):
            if sample_offset + stride_size >= speech_length - 1:
                stride_size = speech_length - sample_offset
                is_final = True

            data = speech[sample_offset: sample_offset + stride_size]
            message = data
            # voices.put(message)
            await websocket.send(message)
            await asyncio.sleep(0.02 * 10)

        await websocket.send(
            '{"command":"END"}')


async def message(id):
    global websocket,voices,offline_msg_done
    text_print = ""
    text_print_2pass_online = ""
    text_print_2pass_offline = ""
    if args.output_dir is not None:
        ibest_writer = open(os.path.join(args.output_dir, "text.{}".format(id)), "a", encoding="utf-8")
    else:
        ibest_writer = None
    try:
       while True:
        
            meg = await websocket.recv()
            meg = json.loads(meg)
            text = meg.get("variable").get("content")
            audios = meg.get("variable").get("audios")
            print(f'recv text:{text}')
            if audios:
                with open('audio.txt', 'w') as file:
                    file.writelines(audios)
                    file.close()


    except Exception as e:
            print("Exception:", e)
            #traceback.print_exc()
            #await websocket.close()
 



async def ws_client(id, chunk_begin, chunk_size):
  chunk_begin=0
  chunk_size=1
  global websocket,voices,offline_msg_done
  headers = {
      "Upgrade": "websocket",
      "Connection": "Upgrade",
      "Sec-WebSocket-Version": "13",
      "Sec-WebSocket-Key": "base64-encoded-secret-key",
      "Authorization": "Access=1"
  }
 
  for i in range(chunk_begin,chunk_begin+chunk_size):
    offline_msg_done=False
    voices = Queue()
    #uri = "ws://{}:{}/ws/open/asr/sentence-stream".format('api.ai.rmd.com.cn', 80)
    #uri = "ws://{}:{}/ws/open/asr/pass-two-stream".format('127.0.0.1', 37091)
    uri = "ws://{}:{}/ws/open/asr/pass-two-stream".format('192.168.0.179', 37091)
    ssl_context = None
    print("connect to", uri)
    async with websockets.connect(uri, ping_interval=None, ssl=ssl_context, extra_headers=headers) as websocket:
        task = asyncio.create_task(record_fs())
        task3 = asyncio.create_task(message(str(id)+"_"+str(i))) #processid+fileid
        await asyncio.gather(task, task3)
  exit(0)
    

def one_thread(id, chunk_begin, chunk_size):
    asyncio.get_event_loop().run_until_complete(ws_client(id, chunk_begin, chunk_size))
    asyncio.get_event_loop().run_forever()

if __name__ == '__main__':
    # for microphone
    p = Process(target=one_thread, args=(0, 0, 0))
    p.start()
    p.join()
    print('end')

